package com.huawei.dli.flink.sye.sink;

import com.huawei.dli.flink.sye.utils.LoginUtil;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HTableDescriptor;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.security.UserGroupInformation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.security.PrivilegedExceptionAction;

public class HBaseSinkAvro extends RichSinkFunction<String> {

  private static final Logger logger = LoggerFactory.getLogger(HBaseSinkAvro.class);

  private org.apache.hadoop.conf.Configuration config = null;
  private Connection conn = null;
  private Table table = null;

  private String tableName = "test_namespace:krb_t1";
  private String familyName = "f1";
  private String columnName = "c1";

  private String userPrincipal;
  private String tmpDir = "/tmp/hbaseSink.";

  public HBaseSinkAvro(String principal) {
    this.userPrincipal = principal;
  }

  @Override
  public void open(Configuration parameters) throws Exception {
    config = HBaseConfiguration.create();
    config.set("hadoop.security.authentication", "kerberos");
    config.set("hadoop.rpc.protection", "privacy");
    config.addResource(this.getClass().getResource("/hbase-site.xml"));
    connect();
  }

  @Override
  public void invoke(String value, Context context) throws Exception {
    try {
      Put put = new Put(Bytes.toBytes(System.currentTimeMillis()));
      put.addColumn(Bytes.toBytes(familyName), Bytes.toBytes(columnName), Bytes.toBytes(value));
      table.put(put);
    } catch (Exception e) {
      logger.error("invoke catch exception: " + e);
      e.printStackTrace();
      resetConnect();
    }
  }

  @Override
  public void close() throws Exception {
    super.close();
    try {
      if (table != null) {
        table.close();
      }
      if (conn != null && !conn.isClosed()) {
        conn.close();
      }
      LoginUtil.deleteFile(tmpDir, "user.keytab");
      LoginUtil.deleteFile(tmpDir, "krb5.conf");
    } catch (Exception e) {
      logger.error("close catch exception: " + e);
    }

  }

  private void connect() throws Exception {
    String userKeytabFile = LoginUtil.createFile(tmpDir, "user.keytab", this.getClass());
    String krb5File = LoginUtil.createFile(tmpDir, "krb5.conf", this.getClass());
    LoginUtil.setJaasConf(LoginUtil.Module.ZOOKEEPER.getName(), userPrincipal, userKeytabFile);
    LoginUtil.login(userPrincipal, userKeytabFile, krb5File, config);

    UserGroupInformation.getLoginUser().doAs(new PrivilegedExceptionAction<Void>() {
      @Override
      public Void run() throws Exception {
        try {
          conn = ConnectionFactory.createConnection(config);
          TableName tableName = TableName.valueOf(HBaseSinkAvro.this.tableName);
          Admin admin = conn.getAdmin();
          if (!admin.tableExists(tableName)) {
            admin.createTable(new HTableDescriptor(tableName).addFamily(new HColumnDescriptor(familyName)));
          }
          table = conn.getTable(tableName);
          logger.info("connect success");
        } catch (Exception e) {
          logger.info("connect failed");
          throw e;
        }
        return null;
      }
    });
  }

  private void resetConnect() throws Exception {
    if (conn != null && !conn.isClosed()) {
      conn.close();
    }
    connect();
    logger.info("reset connect success");
  }
}
